import pandas as pd
import pymysql
from pymysql import IntegrityError
from typing import List, Dict, Optional, Any, Tuple
import numpy as np
from python_depend.gms_api_client import get_quote_kline
from python_depend.mysql_pool import MySQLPool

# ==============================
# 数据库操作函数
# ==============================
def fetch_stock_list() -> List[Dict[str, Any]]:
    mysql_pool = MySQLPool.get_pool('ainvestdb')
    """从数据库获取股票列表"""
    sql_query = """
    SELECT
        CAST(thsmarket_code_hq + 256 AS SIGNED) AS market,
        thscode_hq AS code
    FROM
        pub205
    WHERE
        F005V_PUB205 IN ('212049', '212050', '212010', '212011', '212210')
        AND f018v_pub205 = '正常上市'
        AND F007D_PUB205 IS NOT NULL
        AND thscode_hq IS NOT NULL
        AND F003V_PUB205 IN ('海外普通股', '优先股', '存托凭证(普通股)', '存托凭证(优先股)', 'UNIT');
    """
    try:
        with mysql_pool.connection() as connection:
            with connection.cursor() as cursor:
                cursor.execute(sql_query)
                return cursor.fetchall()
    except pymysql.Error as e:
        print(f"[ERROR] 获取股票列表失败: {e}")
        return []


# ==============================
# 数据写入函数
# ==============================
def write_to_mysql(df: pd.DataFrame, market: str, symbol: str) -> None:
    mysql_pool = MySQLPool.get_pool('ext_quote_data')
    """批量写入MySQL数据库"""
    if df.empty:
        print(f"[INFO] {market}/{symbol} 无有效信号，跳过写入")
        return

    insert_query = """
    INSERT INTO indicator_volume_spikes (
        `market`, `symbol`, `trade_date`, `signal`, 
        `open`, `high`, `low`, `close`, 
        `volume`, `volume_ma`, `volume_ratio`, 
        `create_time`, `update_time`
    ) VALUES (
        %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), NOW()
    )
    ON DUPLICATE KEY UPDATE
        `signal` = VALUES(`signal`),
        `open` = VALUES(`open`),
        `high` = VALUES(`high`),
        `low` = VALUES(`low`),
        `close` = VALUES(`close`),
        `volume` = VALUES(`volume`),
        `volume_ma` = VALUES(`volume_ma`),
        `volume_ratio` = VALUES(`volume_ratio`)
    """

    # 准备批量插入数据
    data = df[['trade_date', 'signal_type', 'open', 'high', 'low', 'close',
               'volume', 'volume_sma', 'volume_ratio']].values
    params = [(market, symbol) + tuple(row) for row in data]

    try:
        connection = mysql_pool.connection()
        with connection.cursor() as cursor:
            cursor.executemany(insert_query, params)
            connection.commit()
            print(f"[SUCCESS] 写入{len(params)}条记录到{market}/{symbol}")
    except IntegrityError as e:
        print(f"[WARNING] {market}/{symbol} 唯一键冲突: {e}")
    except pymysql.Error as e:
        print(f"[ERROR] 写入{market}/{symbol}失败: {e}")



# ==============================
# K线数据处理函数
# ==============================
def fetch_kline_data(market: str, code: str, count: int = 200) -> Optional[pd.DataFrame]:
    """获取K线数据并转换为DataFrame"""

    request_params = {
        "code_list": [{"market": market, "codes": [code]}],
        "trade_class": "intraday",
        "time_period": "day_1",
        "time_range": {"count": count, "end_time": 0},
        "adjust_type": "forward"
    }
    try:
        result = get_quote_kline(request_params)
        return convert_to_dataframe(result)
    except Exception as e:
        print(f"[ERROR] 获取{market}/{code} K线数据失败: {e}")
        return pd.DataFrame


def convert_to_dataframe(data: Dict[str, Any]) -> pd.DataFrame:
    """将API返回数据解析为DataFrame"""
    if data.get('data') is None:
        return pd.DataFrame()
    stock_data = data.get('data', {}).get('quote_data', [{}])[0]
    data_fields = stock_data.get('data_fields', [])
    value_list = stock_data.get('value', [])
    # 字段映射与校验
    field_mapping = {
        'trade_day': '65558',  # 交易日
        'open': '7',
        'high': '8',
        'low': '9',
        'close': '11',
        'volume': '13'
    }

    missing_fields = [f for f in field_mapping.values() if f not in data_fields]
    if missing_fields:
        print(f"[WARNING] 缺失关键字段: {missing_fields}，跳过数据转换")
        return pd.DataFrame()

    # 提取字段索引
    idx_map = {key: data_fields.index(value) for key, value in field_mapping.items()}

    # 构建DataFrame
    try:
        df = pd.DataFrame({
            'open': [float(values[idx_map['open']]) if values[idx_map['open']] else np.nan for values in value_list],
            'high': [float(values[idx_map['high']]) if values[idx_map['high']] else np.nan for values in value_list],
            'low': [float(values[idx_map['low']]) if values[idx_map['low']] else np.nan for values in value_list],
            'close': [float(values[idx_map['close']]) if values[idx_map['close']] else np.nan for values in value_list],
            'volume': [float(values[idx_map['volume']]) if values[idx_map['volume']] else np.nan for values in
                       value_list],
        }, index=[
            values[idx_map['trade_day']] for values in value_list
        ])
        df.index.name = 'trade_day'
        return df.dropna()  # 过滤无效数据
    except Exception as e:
        print(f"[ERROR] 数据转换失败: {e}")
        return pd.DataFrame()


# ==============================
# 指标计算类
# ==============================
class VolumeSpikesIndicator:
    """成交量 spike 指标计算"""

    def __init__(self,
                 vol_multiplier: float = 1.5,
                 vol_ma_length: int = 100,
                 only_valid_hl: bool = True,
                 only_hammers_shooters: bool = True,
                 only_same_color: bool = False):
        self.vol_multiplier = vol_multiplier
        self.vol_ma_length = vol_ma_length
        self.only_valid_hl = only_valid_hl
        self.only_hammers_shooters = only_hammers_shooters
        self.only_same_color = only_same_color

    def _calculate_valid_highs_lows(self, df: pd.DataFrame) -> Tuple[pd.Series, pd.Series]:
        """计算有效高低点"""
        if not self.only_valid_hl:
            return pd.Series(True, index=df.index), pd.Series(True, index=df.index)

        valid_high = (
                (df['high'].shift(1) > df['high'].shift(2)) &
                (df['high'].shift(1) > df['high'])
        ).fillna(False)

        valid_low = (
                (df['low'].shift(1) < df['low'].shift(2)) &
                (df['low'].shift(1) < df['low'])
        ).fillna(False)

        return valid_high, valid_low

    def _calculate_hammer_shooter_patterns(self, df: pd.DataFrame) -> Tuple[pd.Series, pd.Series]:
        """计算锤子/射击之星形态"""
        distance_hl = df['high'] - df['low']
        midpoint = df['low'] + distance_hl / 2

        valid_hammer = (df['open'] > midpoint) & (df['close'] > midpoint)
        valid_shooter = (df['open'] < midpoint) & (df['close'] < midpoint)

        return valid_hammer.fillna(False), valid_shooter.fillna(False)

    def calculate_signals(self, df: pd.DataFrame) -> Dict[str, Any]:
        """主信号计算函数"""
        if len(df) < max(self.vol_ma_length + 2, 3):
            raise ValueError(f"数据不足，至少需要{max(self.vol_ma_length + 2, 3)}条记录")

        df = df.sort_index().copy()
        # 计算有效高低点
        valid_high, valid_low = self._calculate_valid_highs_lows(df)
        valid_hammer, valid_shooter = self._calculate_hammer_shooter_patterns(df)

        if self.only_hammers_shooters:
            valid_high = valid_high & valid_shooter.shift(1)
            valid_low = valid_low & valid_hammer.shift(1)

        if self.only_same_color:
            prev_close = df['close'].shift(1)
            prev_open = df['open'].shift(1)
            valid_high = valid_high & (prev_close < prev_open)  # 前一根阴线
            valid_low = valid_low & (prev_close > prev_open)  # 前一根阳线

        # 验证成交量
        vol_sma = df['volume'].rolling(window=self.vol_ma_length).mean()
        vol_check = df['volume'] > vol_sma * self.vol_multiplier

        result_bearish = valid_high & vol_check.shift(1).fillna(False)
        result_bullish = valid_low & vol_check.shift(1).fillna(False)

        # 实际结果往前偏移一个单位
        result_bearish = result_bearish.shift(-1).fillna(False)
        result_bullish = result_bullish.shift(-1).fillna(False)

        return {
            'bearish_signals': result_bearish,
            'bullish_signals': result_bullish,
            'volume_sma': vol_sma,
            'valid_high': valid_high,
            'valid_low': valid_low,
            'volume_check': vol_check
        }

    def get_signal_summary(self, df: pd.DataFrame) -> pd.DataFrame:
        """生成信号摘要"""
        try:
            results = self.calculate_signals(df)
            summary = []
            # 处理看跌信号
            bearish_indices = df.index[results['bearish_signals']]
            for ts in bearish_indices:
                row = df.loc[ts]
                sma = results['volume_sma'].loc[ts]
                summary.append({
                    'trade_date': ts,
                    'signal_type': 'Bearish',
                    'open': row['open'],
                    'high': row['high'],
                    'low': row['low'],
                    'close': row['close'],
                    'volume': row['volume'],
                    'volume_sma': sma,
                    'volume_ratio': float(f"{row['volume'] / sma:.4f}") if sma != 0 else np.nan
                })

            # 处理看涨信号
            bullish_indices = df.index[results['bullish_signals']]
            for ts in bullish_indices:
                row = df.loc[ts]
                sma = results['volume_sma'].loc[ts]
                summary.append({
                    'trade_date': ts,
                    'signal_type': 'Bullish',
                    'open': row['open'],
                    'high': row['high'],
                    'low': row['low'],
                    'close': row['close'],
                    'volume': row['volume'],
                    'volume_sma': sma,
                    'volume_ratio': float(f"{row['volume'] / sma:.4f}") if sma != 0 else np.nan
                })

            if not summary:
                return pd.DataFrame()

            return pd.DataFrame(summary)
        except Exception as e:
            print(f"[ERROR] 信号计算失败: {e}")
            return pd.DataFrame()


# ==============================
# 主执行函数
# ==============================
def execute_job(vol_multiplier, vol_ma_length, k_count):
    """主任务执行流程"""
    print("[START] 开始执行任务")
    # 1. 获取股票列表
    stock_list = fetch_stock_list()
    # stock_list = [{"code": "AAPL", "market": 185}]
    if not stock_list:
        print("[ERROR] 股票列表为空，任务终止")
        return
    print("查询股票数量：" + str(len(stock_list)))
    # 2. 处理每只股票
    for stock in stock_list:
        market = str(stock['market'])
        code = stock['code']
        print(f"\n[PROCESS] 处理股票: Market={market}, Code={code}")
        # 获取K线数据
        df = fetch_kline_data(market=market, code=code,count=k_count)
        if df.empty:
            print(f"[WARNING] {market}/{code} 无有效K线数据")
            continue
        # 计算信号
        indicator = VolumeSpikesIndicator(vol_multiplier=vol_multiplier,vol_ma_length=vol_ma_length)
        signal_summary = indicator.get_signal_summary(df)
        print(signal_summary)
        # 写入数据库
        if not signal_summary.empty:
            write_to_mysql(signal_summary, market, code)
        else:
            print(f"[INFO] {market}/{code} 无交易信号")

    print("[END] 任务执行完成")


# ==============================
# 程序入口
# ==============================
if __name__ == "__main__":
    start_type = "${start_type}"
    print("start params is： "+start_type)
    MySQLPool.initialize()
    # 首次执行从头开始
    if start_type == "first":
        print("首次执行任务，进行历史数据补正")
        execute_job(1.5 , 100, 200)
    else:
    # 每天执行
        print("每天任务调度")
        execute_job(1.5, 100, 103)
    print("关闭数据库连接")
    MySQLPool.close_all()